import torch
from torch.utils.data import Dataset
import os
import cv2
import xml.etree.ElementTree as ET
import torchvision.transforms as transforms
import numpy as np
import random
from Utils import image
from Config.ConfigTrain import *
class VOCDataSet(Dataset):


    def __init__(self, imgs_path="../DataSet/VOC2007+2012/Train/JPEGImages",
                 annotations_path="../DataSet/VOC2007+2012/Train/Annotations",
                 is_train=True, class_num=Classes,
                 label_smooth_value=0.05, input_size=448, grid_size=64):  # input_size:输入图像的尺度
        self.label_smooth_value = label_smooth_value
        self.class_num = class_num
        self.imgs_name = os.listdir(imgs_path)
        self.input_size = input_size
        self.grid_size = grid_size
        self.is_train = is_train

        self.transform_common = transforms.Compose([
            transforms.ToTensor(),  # height * width * channel -> channel * height * width
            transforms.Normalize(mean=(0.408, 0.448, 0.471), std=(0.242, 0.239, 0.234))  # 归一化后.不容易产生梯度爆炸的问题
        ])
        self.imgs_path = imgs_path
        self.annotations_path = annotations_path
        self.class_dict = {}
        class_index = 0
        """
        读取配置标签
        """
        for class_name in ClassesName:
            self.class_dict[class_name] = class_index
            class_index+=1

    def __getitem__(self, item):

        img_path = os.path.join(self.imgs_path, self.imgs_name[item])
        annotation_path = os.path.join(self.annotations_path, self.imgs_name[item].replace(".jpg", ".xml"))
        img = cv2.imread(img_path)
        tree = ET.parse(annotation_path)
        annotation_xml = tree.getroot()

        objects_xml = annotation_xml.findall("object")
        coords = []

        for object_xml in objects_xml:
            bnd_xml = object_xml.find("bndbox")
            class_name = object_xml.find("name").text
            if class_name not in self.class_dict:  # 不属于我们规定的类
                continue
            xmin = round((float)(bnd_xml.find("xmin").text))
            ymin = round((float)(bnd_xml.find("ymin").text))
            xmax = round((float)(bnd_xml.find("xmax").text))
            ymax = round((float)(bnd_xml.find("ymax").text))
            class_id = self.class_dict[class_name]
            """
            这里解析存储的是5个值，缩放，归一化后的坐标和对应的类别的标签
            """
            coords.append([xmin, ymin, xmax, ymax, class_id])

        coords.sort(key=lambda coord: (coord[2] - coord[0]) * (coord[3] - coord[1]))
        if self.is_train:

            transform_seed = random.randint(0, 4)

            if transform_seed == 0:  # 原图
                img, coords = image.resize_image_with_coords(img, self.input_size, self.input_size, coords)
                img = self.transform_common(img)

            elif transform_seed == 1:  # 缩放+中心裁剪
                img, coords = image.center_crop_with_coords(img, coords)
                img, coords = image.resize_image_with_coords(img, self.input_size, self.input_size, coords)
                img = self.transform_common(img)

            elif transform_seed == 2:  # 平移
                img, coords = image.transplant_with_coords(img, coords)
                img, coords = image.resize_image_with_coords(img, self.input_size, self.input_size, coords)
                img = self.transform_common(img)

            elif transform_seed == 3:  # 明度调整 YOLO在论文中称曝光度为明度
                img, coords = image.resize_image_with_coords(img, self.input_size, self.input_size, coords)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
                H, S, V = cv2.split(img)
                cv2.merge([np.uint8(H), np.uint8(S), np.uint8(V * 1.5)], dst=img)
                cv2.cvtColor(src=img, dst=img, code=cv2.COLOR_HSV2BGR)
                img = self.transform_common(img)

            else:  # 饱和度调整
                img, coords = image.resize_image_with_coords(img, self.input_size, self.input_size, coords)
                H, S, V = cv2.split(img)
                cv2.merge([np.uint8(H), np.uint8(S * 1.5), np.uint8(V)], dst=img)
                cv2.cvtColor(src=img, dst=img, code=cv2.COLOR_HSV2BGR)
                img = self.transform_common(img)

        else:
            img, coords = image.resize_image_with_coords(img, self.input_size, self.input_size, coords)
            img = self.transform_common(img)

        ground_truth = self.encode(coords)

        """
        这里传入的coords是经过图片增强，然后归一化之后的
        之后的话，我们需要经过encode目的是的为了制作方便后期和pred对比的label
        """

        return img,ground_truth


    def __len__(self):
        return len(self.imgs_name)

    def encode(self, coords):

        feature_size = self.input_size // self.grid_size
        ground_truth = np.zeros([feature_size, feature_size, 10 + self.class_num],dtype=float)

        for coord in coords:
            # positive_num = positive_num + 1
            # bounding box归一化
            xmin, ymin, xmax, ymax, class_id = coord

            ground_width = (xmax - xmin)
            ground_height = (ymax - ymin)

            center_x = (xmin + xmax) / 2
            center_y = (ymin + ymax) / 2

            index_row = (int)(center_y * feature_size)
            index_col = (int)(center_x * feature_size)


            ground_box = [center_x * feature_size - index_col, center_y * feature_size - index_row,
                          ground_width, ground_height, 1,
                          round(xmin * self.input_size), round(ymin * self.input_size),
                          round(xmax * self.input_size), round(ymax * self.input_size),
                          round(ground_width * self.input_size * ground_height * self.input_size)
                          ]
            # ground_box.extend(class_list)
            class_ = [0 for _ in range(self.class_num)]
            class_[class_id]=1
            ground_box.extend(class_)

            ground_truth[index_row][index_col] = np.array(ground_box,dtype=float)

        return ground_truth



